/*
 * dhmm.cc
 *
 *  Created on: Sep 19, 2011
 *      Author: tqlong
 */

#include <fstream>
#include "dhmm.h"

void DiscreteHMM::load(const char* fileName)
{
  int n_state, n_observation;
  std::ifstream file(fileName);
  if (!file.good())
  {
    std::cerr << "Cannot open file " << fileName << std::endl;
    exit(1);
  }

  file >> n_state >> n_observation;
  start_ = arma::vec(n_state);
  trans_ = arma::mat(n_state, n_state);           // column wise p(i|j)
  emis_  = arma::mat(n_observation, n_state);     // column wise p(o=k|j)
  for (int i = 0; i < n_state; i++)
    file >> start_(i);
  for (int j = 0; j < n_state; j++)
    for (int i = 0; i < n_state; i++)
      file >> trans_(i, j);
  for (int j = 0; j < n_state; j++)
    for (int k = 0; k < n_observation; k++)
      file >> emis_(k, j);
  file.close();
}

void DiscreteHMM::rand() {
  trans_.randu();
  emis_.randu();
  start_.randu();

  for (int j = 0; j < n_state(); j++)
  {
    trans_.col(j) /= arma::sum(trans_.col(j));
    emis_.col(j) /= arma::sum(emis_.col(j));
  }
  start_ /= arma::sum(start_);
}

void DiscreteHMM::generate(int T, joint_seq_t &seq)
{
  seq.oseq.resize(T);
  seq.qseq.resize(T);
  int q = discrete_sample(start_);
  seq.qseq[0] = q;
  seq.oseq[0] = discrete_sample(emis_.col(q));
  for (int t = 1; t < T; t++)
  {
    q = discrete_sample(trans_.col(q));
    seq.qseq[t] = q;
    seq.oseq[t] = discrete_sample(emis_.col(q));
  }
}

// computation functions
double DiscreteHMM::seq_probability(const sequence_t& o)
{
  compute_alpha(o, var_);
  int T = o.size();
  double s = log(0);
  for (int i = 0; i < n_state(); i++)
    s = logadd(s, var_.alpha(i, T-1));
  return s;
}

void DiscreteHMM::map(joint_seq_t& seq)
{
  sequence_t& o = seq.oseq;
  int T = o.size();

  // viterbi
  var_.alpha = arma::mat(n_state(), T);
  for (int i = 0; i < n_state(); i++)
  {
    var_.alpha(i, 0) = log(start_(i)) + log(emis_(o[0], i));
  }
  for (int t = 1; t < T; t++)
  {
    for (int i = 0; i < n_state(); i++)
    {
      double s = log(0);
      for (int j = 0; j < n_state(); j++)
      {
        double tmp = var_.alpha(j, t-1) + log(trans_(i,j)) + log(emis_(o[t],i));
        if (s < tmp) s = tmp;
      }
      var_.alpha(i, t) = s;
    }
  }

  // trace back
  sequence_t& q = seq.qseq;
  q.resize(T);
  int max_i = 0;
  for (int i = 0; i < n_state(); i++)
    if (var_.alpha(i,T-1) > var_.alpha(max_i, T-1)) max_i = i;

  q[T-1] = max_i;
  for (int t = T-2; t >= 0; t--)
  {
    int next_i = max_i;
    double max = log(0);
    max_i = -1;
    for (int i = 0; i < n_state(); i++)
    {
      double tmp = var_.alpha(i, t) + log(trans_(next_i,i));
      if (tmp > max)
      {
        max = tmp;
        max_i = i;
      }
    }
    q[t] = max_i;
  }
}

double DiscreteHMM::seq_probability_reverse(const sequence_t& o)
{
  compute_beta(o, var_);
  double s = log(0);
  for (int i = 0; i < n_state(); i++)
    s = logadd(s, log(start_(i)) + log(emis_(o[0], i)) + var_.beta(i, 0));
  return s;
}

// training
void DiscreteHMM::train(const vseq_t& data)
{
  var_.ACC_start = arma::vec(n_state());
  var_.ACC_trans = arma::mat(n_state(), n_state());
  var_.ACC_emis = arma::mat(n_observation(), n_state());
  for (int iter = 0; iter < 500; iter++)
  {
    var_.ACC_start.fill(log(0));
    var_.ACC_trans.fill(log(0));
    var_.ACC_emis.fill(log(0));

    double likelihood = 0;
    for (unsigned int n = 0; n < data.size(); n++)
    {
      const sequence_t& o = data[n];
      double log_prob = seq_probability(o); // forward
      compute_beta(o, var_);                // backward
      accumulate(o, -log_prob);
      likelihood += log_prob;
    }
    std::cout << "iter = " << iter << " likelihood = " << likelihood << std::endl;

    reestimate();
  }
}

// print functions
void DiscreteHMM::print_params() const
{
  std::cout << "n_state = " << n_state() << std::endl
            << "n_observation = " << n_observation() << std::endl
            << "start = " << std::endl << start_
            << "trans = " << std::endl << trans_
            << "emis = " << std::endl << emis_;
}

void DiscreteHMM::JoinSequence::print() const
{
  std::cout << "qseq = ";
  for (unsigned int i = 0; i < qseq.size(); i++)
    std::cout << qseq[i] << " ";
  std::cout << std::endl << "oseq = ";
  for (unsigned int i = 0; i < oseq.size(); i++)
    std::cout << oseq[i] << " ";
  std::cout << std::endl;
}

std::ostream& operator<<(std::ostream& os, const Log& x)
{
  os << "exp(" << x.s_ << ") = " << exp(x.s_);
  return os;
}

// protected members;
void DiscreteHMM::compute_alpha(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.alpha = arma::mat(n_state(), T);
  for (int i = 0; i < n_state(); i++)
  {
    var.alpha(i, 0) = log(start_(i)) + log(emis_(o[0], i));
  }
  for (int t = 1; t < T; t++)
  {
    for (int i = 0; i < n_state(); i++)
    {
      double s = log(0);
      for (int j = 0; j < n_state(); j++)
      {
        s = logadd(s, var.alpha(j, t-1) + log(trans_(i,j)) + log(emis_(o[t],i)));
      }
      var.alpha(i, t) = s;
    }
  }
}

void DiscreteHMM::compute_beta(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.beta = arma::mat(n_state(), T);
  for (int i = 0; i < n_state(); i++)
  {
    var.beta(i, T-1) = log(1);
  }
  for (int t = T-2; t >=0 ; t--)
  {
    for (int j = 0; j < n_state(); j++)
    {
      double s = log(0);
      for (int i = 0; i < n_state(); i++)
      {
        s = logadd(s, var.beta(i, t+1) + log(trans_(i,j)) + log(emis_(o[t+1],i)));
      }
      var.beta(j, t) = s;
    }
  }
}

void DiscreteHMM::accumulate(const sequence_t& o, double weight)
{
  int T = var_.alpha.n_cols;

  for (int j = 0; j < n_state(); j++)
    for (int i = 0; i < n_state(); i++)
    {
      double s = log(0);
      for (int t = 0; t < T-1; t++)
      {
        double tmp = var_.alpha(j,t) + log(trans_(i,j)) + log(emis_(o[t+1], i)) + var_.beta(i,t+1);
        s = logadd(s, tmp);
      }
      s += weight;
      var_.ACC_trans(i,j) = logadd(var_.ACC_trans(i,j), s);
    }

  for (int t = 0; t < T; t++)
  {
    int k = o[t];
    for (int i = 0; i < n_state(); i++)
    {
      double s = var_.alpha(i,t) + var_.beta(i, t) + weight;
      var_.ACC_emis(k, i) = logadd(var_.ACC_emis(k, i), s);
    }
  }

  for (int i = 0; i < n_state(); i++)
  {
    double s = var_.alpha(i,0) + var_.beta(i,0) + weight;
    var_.ACC_start(i) = logadd(var_.ACC_start(i), s);
  }
}

void DiscreteHMM::reestimate()
{
  for (int j = 0; j < n_state(); j++)
  {
    double s = logsum_vec(var_.ACC_trans.col(j));
    for (int i = 0; i < n_state(); i++)
      trans_(i,j) = exp(var_.ACC_trans(i, j) - s);
  }

  for (int j = 0; j < n_state(); j++)
  {
    double s = logsum_vec(var_.ACC_emis.col(j));
    for (int k = 0; k < n_observation(); k++)
      emis_(k,j) = exp(var_.ACC_emis(k, j) - s);
  }

  double s = logsum_vec(var_.ACC_start);
  for (int i = 0; i < n_state(); i++)
    start_(i) = exp(var_.ACC_start(i) - s);
}

// common used functions
int discrete_sample(arma::vec p)
{
  double s = arma::sum(p);
  double r = arma::randu() * s;
  s = 0;
  for (unsigned int i = 0; i < p.n_elem; i++)
  {
    s += p(i);
    if (s >= r) return i;
  }
  return p.n_elem-1;
}

/** Log space computation */
double logadd(double a, double b)
{
  if (a == -std::numeric_limits<double>::infinity())
    return b;
  else if (b == -std::numeric_limits<double>::infinity())
    return a;
  else if (a < b)
    return b+log1p(exp(a-b));
  else // a >= b
    return a+log1p(exp(b-a));
}

double logsub(double a, double b)
{
  if (a == -std::numeric_limits<double>::infinity())
    return log(0);
  else if (a < b)
    return log(0);
  else if (b == -std::numeric_limits<double>::infinity())
    return a;
  else
    return a+log1p(-exp(b-a));
}

double logsum_vec(arma::vec v)
{
  double s = log(0);
  for (unsigned int i = 0; i < v.n_elem; i++)
    s = logadd(s, v.at(i));
  return s;
}

double logsum(arma::mat v)
{
  double s = log(0);
  for (unsigned int i = 0; i < v.n_rows; i++)
    for (unsigned int j = 0; j < v.n_cols; j++)
      s = logadd(s, v.at(i,j));
  return s;
}
